"""
此框架会随机死机，但是执行速度不错
    只能说尝试性质写写，当作onnx完成后的提升性质去实现
"""

import time
import cv2
import numpy as np
import ncnn

net = ncnn.Net()
net.load_param("../2-ReplaceWithSimplePy/v1_ncnn.param")
net.load_model("../2-ReplaceWithSimplePy/v1_ncnn.bin")
ex = net.create_extractor()

def torch_style_softmax_numpy_impl(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

def ncnn_infer_data(mat):
    out = []

    # with ncnn.Net() as net:
    #     net.load_param("v1_onnx.ncnn.param")
    #     net.load_model("v1_onnx.ncnn.bin")
    #
    #     with net.create_extractor() as ex:
    #         # ex.input("in0", ncnn.Mat(in0.squeeze(0).numpy()).clone())
    #         ex.input("in0", ncnn.Mat(mat[0]).clone())
    #
    #         _, out0 = ex.extract("out0")
    #         # out.append(torch.from_numpy(np.array(out0)).unsqueeze(0))
    #         out.append(np.expand_dims(np.array(out0), 0))

    ex.input("input:0", ncnn.Mat(mat[0]))
    _, out0 = ex.extract("output:0")
    out.append(np.expand_dims(np.array(out0), 0))

    if len(out) == 1:
        return out[0]
    else:
        return tuple(out)


def get_nnunet_predict(x: np.ndarray, dims: list):
    return np.flip(
        torch_style_softmax_numpy_impl(ncnn_infer_data(
            np.flip(x, dims)
        ), 1),
        dims
    )

if __name__ == "__main__":
    sample_image = "../resource/raw_resize/1.jpg"
    start_time = time.perf_counter()
    mat = cv2.imread(sample_image)
    mat = mat.transpose((2, 0, 1))
    mat = mat.astype(np.float32)
    for c in range(mat.shape[0]):
        mat[c] = (mat[c] - mat[c].mean()) / (mat[c].std() + 1e-8)
    mat = np.expand_dims(mat, 0)

    pred = torch_style_softmax_numpy_impl(ncnn_infer_data(mat), 1)
    pred += get_nnunet_predict(mat, [3])
    pred += get_nnunet_predict(mat, [2])
    pred += get_nnunet_predict(mat, [3, 2])
    pred /= 4.0

    # save
    pred = np.argmax(pred, 1)
    pred = pred[0]
    cv2.imwrite("out_pred_ncnn.png", pred * 255)
    end_time = time.perf_counter()
    print("time cost: ", end_time - start_time)